import sys
import json
import re
import openpyxl
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QLineEdit, QPushButton, QComboBox, QTableWidget,
    QTableWidgetItem, QFileDialog, QMessageBox, QSpinBox, QGroupBox,
    QRadioButton, QTabWidget, QPlainTextEdit, QSizePolicy
)
from PyQt6.QtCore import Qt

class ExcelExtractorApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Excel单元格提取工具")
        self.setGeometry(100, 50, 920, 620)
        
        # 主控件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.layout = QVBoxLayout()
        self.layout.setContentsMargins(10, 10, 10, 10)
        self.layout.setSpacing(10)
        
        # 文件选择区域
        self.setup_file_selection()
        
        # 工作表选择区域
        self.setup_sheet_selection()
        
        # 字段配置表格
        self.setup_mapping_table()
        
        # 操作按钮（保持原始位置）
        self.setup_buttons()
        
        self.central_widget.setLayout(self.layout)
    
    def setup_file_selection(self):
        """设置文件选择区域"""
        file_layout = QHBoxLayout()
        file_layout.setSpacing(10)
        
        # 输入文件
        self.input_file_edit = QLineEdit()
        self.input_file_edit.setPlaceholderText("选择输入Excel文件...")
        input_btn = QPushButton("浏览...")
        input_btn.setFixedWidth(80)
        input_btn.clicked.connect(self.select_input_file)
        
        # 输出文件
        self.output_file_edit = QLineEdit()
        self.output_file_edit.setPlaceholderText("设置输出Excel路径...")
        output_btn = QPushButton("浏览...")
        output_btn.setFixedWidth(80)
        output_btn.clicked.connect(self.select_output_file)
        
        file_layout.addWidget(QLabel("输入文件:"), stretch=0)
        file_layout.addWidget(self.input_file_edit, stretch=2)
        file_layout.addWidget(input_btn, stretch=0)
        file_layout.addWidget(QLabel("输出文件:"), stretch=0)
        file_layout.addWidget(self.output_file_edit, stretch=2)
        file_layout.addWidget(output_btn, stretch=0)
        
        self.layout.addLayout(file_layout)
    
    def setup_sheet_selection(self):
        """设置工作表选择区域（调整大小）"""
        sheet_group = QGroupBox("工作表选择")
        sheet_layout = QVBoxLayout()
        sheet_group.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed)
        sheet_group.setFixedHeight(280)
        
        self.sheet_tabs = QTabWidget()
        
        # 基本选择方式选项卡
        basic_tab = QWidget()
        basic_layout = QVBoxLayout()
        basic_layout.setSpacing(8)
        basic_layout.setContentsMargins(8, 8, 8, 8)
        
        self.all_sheets_radio = QRadioButton("全部工作表")
        self.all_sheets_radio.setChecked(True)
        basic_layout.addWidget(self.all_sheets_radio)
        
        first_n_layout = QHBoxLayout()
        self.first_n_radio = QRadioButton("前")
        self.first_n_spin = QSpinBox()
        self.first_n_spin.setRange(1, 1000)
        self.first_n_spin.setValue(1)
        self.first_n_spin.setEnabled(False)
        first_n_layout.addWidget(self.first_n_radio)
        first_n_layout.addWidget(self.first_n_spin)
        first_n_layout.addWidget(QLabel("个工作表"))
        first_n_layout.addStretch()
        basic_layout.addLayout(first_n_layout)
        
        last_n_layout = QHBoxLayout()
        self.last_n_radio = QRadioButton("后")
        self.last_n_spin = QSpinBox()
        self.last_n_spin.setRange(1, 1000)
        self.last_n_spin.setValue(1)
        self.last_n_spin.setEnabled(False)
        last_n_layout.addWidget(self.last_n_radio)
        last_n_layout.addWidget(self.last_n_spin)
        last_n_layout.addWidget(QLabel("个工作表"))
        last_n_layout.addStretch()
        basic_layout.addLayout(last_n_layout)
        
        range_layout = QHBoxLayout()
        self.range_radio = QRadioButton("从第")
        self.start_spin = QSpinBox()
        self.start_spin.setRange(1, 1000)
        self.start_spin.setValue(1)
        self.start_spin.setEnabled(False)
        self.end_spin = QSpinBox()
        self.end_spin.setRange(1, 1000)
        self.end_spin.setValue(1)
        self.end_spin.setEnabled(False)
        range_layout.addWidget(self.range_radio)
        range_layout.addWidget(self.start_spin)
        range_layout.addWidget(QLabel("到第"))
        range_layout.addWidget(self.end_spin)
        range_layout.addWidget(QLabel("个工作表"))
        range_layout.addStretch()
        basic_layout.addLayout(range_layout)
        
        basic_tab.setLayout(basic_layout)
        self.sheet_tabs.addTab(basic_tab, "基本选择")
        
        # 正则表达式选择选项卡
        regex_tab = QWidget()
        regex_layout = QVBoxLayout()
        regex_layout.setSpacing(8)
        regex_layout.setContentsMargins(8, 8, 8, 8)
        
        self.regex_radio = QRadioButton("使用正则表达式匹配工作表名")
        regex_layout.addWidget(self.regex_radio)
        
        regex_input_layout = QHBoxLayout()
        regex_input_layout.addWidget(QLabel("正则表达式:"))
        self.regex_edit = QLineEdit()
        self.regex_edit.setPlaceholderText(r"例如: ^Sheet\d+$ 或 202\d_.*")
        regex_input_layout.addWidget(self.regex_edit)
        regex_layout.addLayout(regex_input_layout)
        
        self.regex_test_edit = QPlainTextEdit()
        self.regex_test_edit.setPlaceholderText("匹配到的工作表将显示在这里...")
        self.regex_test_edit.setReadOnly(True)
        self.regex_test_edit.setMaximumHeight(80)
        regex_layout.addWidget(QLabel("匹配结果:"))
        regex_layout.addWidget(self.regex_test_edit)
        
        test_btn = QPushButton("测试正则表达式")
        test_btn.setFixedWidth(120)
        regex_layout.addWidget(test_btn, alignment=Qt.AlignmentFlag.AlignRight)
        test_btn.clicked.connect(self.test_regex_pattern)
        
        regex_tab.setLayout(regex_layout)
        self.sheet_tabs.addTab(regex_tab, "正则匹配")
        
        sheet_layout.addWidget(self.sheet_tabs)
        
        # 连接信号
        self.all_sheets_radio.toggled.connect(self.update_sheet_selection)
        self.first_n_radio.toggled.connect(self.update_sheet_selection)
        self.last_n_radio.toggled.connect(self.update_sheet_selection)
        self.range_radio.toggled.connect(self.update_sheet_selection)
        self.regex_radio.toggled.connect(self.update_sheet_selection)
        
        sheet_group.setLayout(sheet_layout)
        self.layout.addWidget(sheet_group)
    
    def setup_mapping_table(self):
        """设置字段配置表格（简化版，只保留列名和单元格引用）"""
        self.table = QTableWidget()
        self.table.setColumnCount(2)  # 只保留两列
        self.table.setHorizontalHeaderLabels(["列名", "单元格引用"])
        
        # 设置列宽
        self.table.setColumnWidth(0, 300)
        self.table.setColumnWidth(1, 150)
        
        # 表格属性
        self.table.horizontalHeader().setStretchLastSection(True)
        self.table.verticalHeader().setDefaultSectionSize(35)
        self.table.setSelectionBehavior(QTableWidget.SelectionBehavior.SelectRows)
        self.table.setSelectionMode(QTableWidget.SelectionMode.SingleSelection)
        self.table.setMinimumHeight(200)
        
        # 添加3行初始数据
        for _ in range(3):
            self.add_mapping_row()
        
        config_label = QLabel("字段配置:")
        config_label.setStyleSheet("font-weight: bold;")
        self.layout.addWidget(config_label)
        self.layout.addWidget(self.table)
    
    def setup_buttons(self):
        """设置操作按钮（添加行顺序调整功能）"""
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(10)
        
        add_btn = QPushButton("添加字段")
        add_btn.setFixedWidth(100)
        add_btn.clicked.connect(self.add_mapping_row)
        
        remove_btn = QPushButton("删除选中")
        remove_btn.setFixedWidth(100)
        remove_btn.clicked.connect(self.remove_selected_rows)
        
        # 添加行顺序调整按钮
        move_up_btn = QPushButton("上移")
        move_up_btn.setFixedWidth(80)
        move_up_btn.clicked.connect(self.move_row_up)
        
        move_down_btn = QPushButton("下移")
        move_down_btn.setFixedWidth(80)
        move_down_btn.clicked.connect(self.move_row_down)
        
        save_btn = QPushButton("保存配置")
        save_btn.setFixedWidth(100)
        save_btn.clicked.connect(self.save_config)
        
        load_btn = QPushButton("加载配置")
        load_btn.setFixedWidth(100)
        load_btn.clicked.connect(self.load_config)
        
        self.run_btn = QPushButton("执行提取")
        self.run_btn.setFixedWidth(100)
        self.run_btn.setStyleSheet("background-color: #4CAF50; color: white;")
        self.run_btn.clicked.connect(self.run_extraction)
        
        btn_layout.addWidget(add_btn)
        btn_layout.addWidget(remove_btn)
        btn_layout.addWidget(move_up_btn)  # 添加上移按钮
        btn_layout.addWidget(move_down_btn)  # 添加下移按钮
        btn_layout.addWidget(save_btn)
        btn_layout.addWidget(load_btn)
        btn_layout.addStretch()
        btn_layout.addWidget(self.run_btn)
        
        self.layout.addLayout(btn_layout)
    
    def move_row_up(self):
        """将选中的行向上移动"""
        selected_rows = sorted([row.row() for row in self.table.selectionModel().selectedRows()])
        if not selected_rows:
            return
            
        for row in selected_rows:
            if row > 0:  # 如果不是第一行
                self.swap_rows(row, row-1)
                self.table.selectRow(row-1)  # 保持选中状态
    
    def move_row_down(self):
        """将选中的行向下移动"""
        selected_rows = sorted([row.row() for row in self.table.selectionModel().selectedRows()], reverse=True)
        if not selected_rows:
            return
            
        for row in selected_rows:
            if row < self.table.rowCount() - 1:  # 如果不是最后一行
                self.swap_rows(row, row+1)
                self.table.selectRow(row+1)  # 保持选中状态
    
    def swap_rows(self, row1, row2):
        """交换两行的内容"""
        # 保存行1的数据
        row1_data = [
            self.table.cellWidget(row1, 0).text(),
            self.table.cellWidget(row1, 1).text()
        ]
        
        # 保存行2的数据
        row2_data = [
            self.table.cellWidget(row2, 0).text(),
            self.table.cellWidget(row2, 1).text()
        ]
        
        # 交换数据
        self.table.cellWidget(row1, 0).setText(row2_data[0])
        self.table.cellWidget(row1, 1).setText(row2_data[1])
        
        self.table.cellWidget(row2, 0).setText(row1_data[0])
        self.table.cellWidget(row2, 1).setText(row1_data[1])
    
    def select_input_file(self):
        """选择输入文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择Excel文件", "", "Excel文件 (*.xlsx *.xlsm)"
        )
        if file_path:
            self.input_file_edit.setText(file_path)
    
    def select_output_file(self):
        """选择输出文件"""
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存Excel文件", "", "Excel文件 (*.xlsx)"
        )
        if file_path:
            if not file_path.lower().endswith('.xlsx'):
                file_path += '.xlsx'
            self.output_file_edit.setText(file_path)
    
    def add_mapping_row(self, data=None):
        """添加新的配置行（简化版）"""
        row = self.table.rowCount()
        self.table.insertRow(row)
        
        # 列名
        col_name = QLineEdit()
        if data:
            col_name.setText(data.get("column_name", ""))
        self.table.setCellWidget(row, 0, col_name)
        
        # 单元格引用
        cell_ref = QLineEdit()
        if data:
            cell_ref.setText(data.get("cell_reference", ""))
        cell_ref.setPlaceholderText("如: A1")
        self.table.setCellWidget(row, 1, cell_ref)
    
    def remove_selected_rows(self):
        """删除选中行"""
        selected = self.table.selectionModel().selectedRows()
        for idx in sorted([row.row() for row in selected], reverse=True):
            self.table.removeRow(idx)
    
    def update_sheet_selection(self):
        """更新工作表选择控件的启用状态"""
        self.first_n_spin.setEnabled(self.first_n_radio.isChecked())
        self.last_n_spin.setEnabled(self.last_n_radio.isChecked())
        self.start_spin.setEnabled(self.range_radio.isChecked())
        self.end_spin.setEnabled(self.range_radio.isChecked())
        self.regex_edit.setEnabled(self.regex_radio.isChecked())
    
    def test_regex_pattern(self):
        """测试正则表达式模式"""
        if not self.input_file_edit.text():
            QMessageBox.warning(self, "警告", "请先选择输入文件")
            return
            
        pattern = self.regex_edit.text().strip()
        if not pattern:
            QMessageBox.warning(self, "警告", "请输入正则表达式")
            return
            
        try:
            re.compile(pattern)
        except re.error as e:
            QMessageBox.warning(self, "错误", f"无效的正则表达式:\n{str(e)}")
            return
            
        try:
            wb = openpyxl.load_workbook(self.input_file_edit.text(), read_only=True)
            sheet_names = wb.sheetnames
            matched_sheets = [name for name in sheet_names if re.search(pattern, name)]
            
            self.regex_test_edit.setPlainText("\n".join(matched_sheets) if matched_sheets else "没有匹配到任何工作表")
            
            QMessageBox.information(
                self, "测试结果",
                f"匹配到 {len(matched_sheets)} 个工作表\n\n" +
                "\n".join(matched_sheets[:10]) + 
                ("\n..." if len(matched_sheets) > 10 else "")
            )
        except Exception as e:
            QMessageBox.critical(self, "错误", f"测试失败:\n{str(e)}")
    
    def validate_inputs(self):
        """验证输入"""
        if not self.input_file_edit.text():
            QMessageBox.warning(self, "警告", "请选择输入文件")
            return False
        
        if not self.output_file_edit.text():
            QMessageBox.warning(self, "警告", "请设置输出文件")
            return False
        
        if self.regex_radio.isChecked() and not self.regex_edit.text().strip():
            QMessageBox.warning(self, "警告", "请输入正则表达式")
            return False
        
        for row in range(self.table.rowCount()):
            col_name = self.table.cellWidget(row, 0).text().strip()
            cell_ref = self.table.cellWidget(row, 1).text().strip().upper()
            
            if not col_name:
                QMessageBox.warning(self, "警告", f"第{row+1}行: 列名不能为空")
                return False
                
            if not re.match(r'^[A-Z]{1,3}\d{1,7}$', cell_ref):
                QMessageBox.warning(
                    self, "警告", 
                    f"第{row+1}行: 单元格引用格式错误 (应为如A1/BB23)"
                )
                return False
                
        return True
    
    def get_config_data(self):
        """获取当前配置数据（简化版）"""
        config = {
            "input_file": self.input_file_edit.text(),
            "output_file": self.output_file_edit.text(),
            "sheet_selection": {
                "all_sheets": self.all_sheets_radio.isChecked(),
                "first_n": self.first_n_spin.value() if self.first_n_radio.isChecked() else 0,
                "last_n": self.last_n_spin.value() if self.last_n_radio.isChecked() else 0,
                "range_start": self.start_spin.value() if self.range_radio.isChecked() else 0,
                "range_end": self.end_spin.value() if self.range_radio.isChecked() else 0,
                "regex_pattern": self.regex_edit.text() if self.regex_radio.isChecked() else ""
            },
            "mappings": []
        }
        
        for row in range(self.table.rowCount()):
            mapping = {
                "column_name": self.table.cellWidget(row, 0).text().strip(),
                "cell_reference": self.table.cellWidget(row, 1).text().strip().upper()
            }
            config["mappings"].append(mapping)
        
        return config
    
    def save_config(self):
        """保存配置到文件"""
        if not self.validate_inputs():
            return
            
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存配置", "", "JSON文件 (*.json)"
        )
        if file_path:
            if not file_path.lower().endswith('.json'):
                file_path += '.json'
                
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(self.get_config_data(), f, indent=2)
                
            QMessageBox.information(self, "成功", f"配置已保存到:\n{file_path}")
    
    def load_config(self):
        """从文件加载配置"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "加载配置", "", "JSON文件 (*.json)"
        )
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                
                while self.table.rowCount() > 0:
                    self.table.removeRow(0)
                
                self.input_file_edit.setText(config.get("input_file", ""))
                self.output_file_edit.setText(config.get("output_file", ""))
                
                sheet_sel = config.get("sheet_selection", {})
                if sheet_sel.get("all_sheets", True):
                    self.all_sheets_radio.setChecked(True)
                elif sheet_sel.get("first_n", 0) > 0:
                    self.first_n_radio.setChecked(True)
                    self.first_n_spin.setValue(sheet_sel["first_n"])
                elif sheet_sel.get("last_n", 0) > 0:
                    self.last_n_radio.setChecked(True)
                    self.last_n_spin.setValue(sheet_sel["last_n"])
                elif sheet_sel.get("range_start", 0) > 0:
                    self.range_radio.setChecked(True)
                    self.start_spin.setValue(sheet_sel["range_start"])
                    self.end_spin.setValue(sheet_sel["range_end"])
                elif sheet_sel.get("regex_pattern", ""):
                    self.regex_radio.setChecked(True)
                    self.regex_edit.setText(sheet_sel["regex_pattern"])
                
                for mapping in config.get("mappings", []):
                    self.add_mapping_row(mapping)
                    
                QMessageBox.information(self, "成功", "配置加载完成")
                
            except Exception as e:
                QMessageBox.critical(self, "错误", f"加载配置失败:\n{str(e)}")
    
    def get_selected_sheets(self, workbook):
        """根据选择获取要处理的工作表列表"""
        sheet_names = workbook.sheetnames
        
        if self.all_sheets_radio.isChecked():
            return sheet_names
        elif self.first_n_radio.isChecked():
            n = min(self.first_n_spin.value(), len(sheet_names))
            return sheet_names[:n]
        elif self.last_n_radio.isChecked():
            n = min(self.last_n_spin.value(), len(sheet_names))
            return sheet_names[-n:]
        elif self.range_radio.isChecked():
            start = max(1, min(self.start_spin.value(), len(sheet_names))) - 1
            end = max(1, min(self.end_spin.value(), len(sheet_names)))
            return sheet_names[start:end]
        elif self.regex_radio.isChecked():
            pattern = self.regex_edit.text().strip()
            try:
                return [name for name in sheet_names if re.search(pattern, name)]
            except re.error:
                return sheet_names
        else:
            return sheet_names
    
    def run_extraction(self):
        """执行数据提取"""
        if not self.validate_inputs():
            return
            
        self.run_btn.setEnabled(False)
        self.run_btn.setText("正在提取...")
        QApplication.processEvents()
        
        try:
            config = self.get_config_data()
            
            input_wb = openpyxl.load_workbook(config["input_file"], data_only=True)
            output_wb = openpyxl.Workbook()
            output_sheet = output_wb.active
            output_sheet.title = "数据汇总"
            
            headers = ["Sheet"] + [m["column_name"] for m in config["mappings"]]
            output_sheet.append(headers)
            
            selected_sheets = self.get_selected_sheets(input_wb)
            
            if not selected_sheets:
                QMessageBox.warning(self, "警告", "没有找到匹配的工作表")
                return
            
            for sheet_name in selected_sheets:
                sheet = input_wb[sheet_name]
                row_data = [sheet_name]
                
                for mapping in config["mappings"]:
                    cell = sheet[mapping["cell_reference"]]
                    value = cell.value if cell.value is not None else ""
                    row_data.append(value)
                
                output_sheet.append(row_data)
            
            output_wb.save(config["output_file"])
            
            QMessageBox.information(
                self, "完成", 
                f"数据提取成功!\n已处理工作表: {len(selected_sheets)}个\n已保存到:\n{config['output_file']}"
            )
            
        except Exception as e:
            QMessageBox.critical(
                self, "错误", 
                f"提取过程中发生错误:\n{str(e)}"
            )
        finally:
            self.run_btn.setEnabled(True)
            self.run_btn.setText("执行提取")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ExcelExtractorApp()
    window.show()
    sys.exit(app.exec())